import string

letters = string.ascii_uppercase

rename_chains = {i: letters[i] for i in range(26)}  # NOTE -- expect errors if you have more than 26 structures
skip_res_list = [
    "HOH",
    "GOL",
    "EDO",
    "SO4",
    "EDO",
    "NAG",
    "PO4",
    "ACT",
    "PEG",
    "MAN",
    "BMA",
    "DMS",
    "MPD",
    "MES",
    "PG4",
    "TRS",
    "FMT",
    "PGE",
    "EPE",
    "NO3",
    "UNX",
    "UNL",
    "UNK",
    "IPA",
    "IMD",
    "GLC",
    "MLI",
    "1PE",
    "NO3",
    "SCN",
    "P6G",
    "OXY",
    "EOH",
    "NH4",
    "DTT",
    "BEN",
    "BCT",
    "FUL",
    "AZI",
    "DOD",
    "OH",
    "CYN",
    "NO",
    "NO2",
    "SO3",
    "H2S",
    "MOH",
    "URE",
    "CO2",
    "2NO",
]  # top ions, sugars, small molecules with N/C/O/S/P that appear to be crystal artifacts or common surface bound co-factors -- to ignore
skip_atoms = ["H", "D"]
atoms = ["N", "C", "O", "S", "P", "other"]
aa = [
    "ALA",
    "ARG",
    "ASN",
    "ASP",
    "CYS",
    "GLN",
    "GLU",
    "GLY",
    "HIS",
    "ILE",
    "LEU",
    "LYS",
    "MET",
    "PHE",
    "PRO",
    "SER",
    "THR",
    "TRP",
    "TYR",
    "VAL",
    "MSE",
]
res_label_dict = {
    "HIS": 0,
    "LYS": 1,
    "ARG": 2,
    "ASP": 3,
    "GLU": 4,
    "SER": 5,
    "THR": 6,
    "ASN": 7,
    "GLN": 8,
    "ALA": 9,
    "VAL": 10,
    "LEU": 11,
    "ILE": 12,
    "MET": 13,
    "PHE": 14,
    "TYR": 15,
    "TRP": 16,
    "PRO": 17,
    "GLY": 18,
    "CYS": 19,
    "MSE": 13,
}  # MSE -- same as MET
label_res_dict = {
    0: "HIS",
    1: "LYS",
    2: "ARG",
    3: "ASP",
    4: "GLU",
    5: "SER",
    6: "THR",
    7: "ASN",
    8: "GLN",
    9: "ALA",
    10: "VAL",
    11: "LEU",
    12: "ILE",
    13: "MET",
    14: "PHE",
    15: "TYR",
    16: "TRP",
    17: "PRO",
    18: "GLY",
    19: "CYS",
}  # , 20:'MSE'}

chi_dict = {
    "ARG": {"chi_1": "CG", "chi_2": "CD", "chi_3": "NE", "chi_4": "CZ"},
    "LYS": {"chi_1": "CG", "chi_2": "CD", "chi_3": "CE", "chi_4": "NZ"},
    "GLN": {"chi_1": "CG", "chi_2": "CD", "chi_3": "OE1"},
    "GLU": {"chi_1": "CG", "chi_2": "CD", "chi_3": "OE1"},
    "MET": {"chi_1": "CG", "chi_2": "SD", "chi_3": "CE"},
    "ASP": {"chi_1": "CG", "chi_2": "OD1"},
    "ILE": {"chi_1": "CG1", "chi_2": "CD1"},
    "HIS": {"chi_1": "CG", "chi_2": "ND1"},
    "LEU": {"chi_1": "CG", "chi_2": "CD1"},
    "ASN": {"chi_1": "CG", "chi_2": "OD1"},
    "PHE": {"chi_1": "CG", "chi_2": "CD1"},
    "PRO": {"chi_1": "CG", "chi_2": "CD"},
    "TRP": {"chi_1": "CG", "chi_2": "CD1"},
    "TYR": {"chi_1": "CG", "chi_2": "CD1"},
    "VAL": {"chi_1": "CG1"},
    "THR": {"chi_1": "OG1"},
    "SER": {"chi_1": "OG"},
    "CYS": {"chi_1": "SG"},
    "GLY": {},
    "ALA": {},
}


chi_dict_old = {
    "ARG": 4,
    "LYS": 4,
    "GLN": 3,
    "GLU": 3,
    "MET": 3,
    "ASP": 2,
    "ILE": 2,
    "HIS": 2,
    "LEU": 2,
    "ASN": 2,
    "PHE": 2,
    "PRO": 3,
    "TRP": 2,
    "TYR": 2,
    "VAL": 1,
    "THR": 1,
    "SER": 1,
    "CYS": 1,
    "GLY": 0,
    "ALA": 0,
}
aa_map = {
    0: "H",
    1: "K",
    2: "R",
    3: "D",
    4: "E",
    5: "S",
    6: "T",
    7: "N",
    8: "Q",
    9: "A",
    10: "V",
    11: "L",
    12: "I",
    13: "M",
    14: "F",
    15: "Y",
    16: "W",
    17: "P",
    18: "G",
    19: "C",
}  # , 20: "M"} # caution methionine in place of MSE
aa_inv = {
    "H": "HIS",
    "K": "LYS",
    "R": "ARG",
    "D": "ASP",
    "E": "GLU",
    "S": "SER",
    "T": "THR",
    "N": "ASN",
    "Q": "GLN",
    "A": "ALA",
    "V": "VAL",
    "L": "LEU",
    "I": "ILE",
    "M": "MET",
    "F": "PHE",
    "Y": "TYR",
    "W": "TRP",
    "P": "PRO",
    "G": "GLY",
    "C": "CYS",
}
aa_map_inv = {
    "H": 0,
    "K": 1,
    "R": 2,
    "D": 3,
    "E": 4,
    "S": 5,
    "T": 6,
    "N": 7,
    "Q": 8,
    "A": 9,
    "V": 10,
    "L": 11,
    "I": 12,
    "M": 13,
    "F": 14,
    "Y": 15,
    "W": 16,
    "P": 17,
    "G": 18,
    "C": 19,
}
aa_to_letter = {aa_inv[k]: k for k in aa_inv.keys()}
label_res_single_dict = {
    0: "H",
    1: "K",
    2: "R",
    3: "D",
    4: "E",
    5: "S",
    6: "T",
    7: "N",
    8: "Q",
    9: "A",
    10: "V",
    11: "L",
    12: "I",
    13: "M",
    14: "F",
    15: "Y",
    16: "W",
    17: "P",
    18: "G",
    19: "C",
}
# resfile commands where values are amino acids allowed by that command 
resfile_commands = {
    "ALLAA": {'H', 'K', 'R', 'D', 'E', 'S', 'T', 'N', 'Q', 'A', 'V', 'L', 'I', 'M', 'F', 'Y', 'W', 'P', 'G', 'C'},
    "ALLAAwc": {'H', 'K', 'R', 'D', 'E', 'S', 'T', 'N', 'Q', 'A', 'V', 'L', 'I', 'M', 'F', 'Y', 'W', 'P', 'G', 'C'},
    "ALLAAxc": {'H', 'K', 'R', 'D', 'E', 'S', 'T', 'N', 'Q', 'A', 'V', 'L', 'I', 'M', 'F', 'Y', 'W', 'P', 'G'},
    "POLAR": {'E', 'H', 'K', 'N', 'R', 'Q', 'D', 'S', 'T'},
    "APOLAR": {'P', 'M', 'Y', 'V', 'F', 'L', 'I', 'A', 'C', 'W', 'G'},
}
